import pandas as pd
import nibabel as nib
import numpy as np
from pathlib import Path
from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,
                             QLineEdit, QPushButton, QTextEdit, QComboBox, QListWidget,
                             QLabel, QSizePolicy)
from PyQt5.QtCore import Qt, QStringListModel
from PyQt5.QtGui import QFont, QKeyEvent, QFocusEvent
import sys


# --------------------------
# 1. 基础配置与数据加载函数（无修改，确保后台逻辑稳定）
# --------------------------
class BrainConnectionTool:
    def __init__(self):
        # 1. 路径初始化
        self.project_root = Path.cwd()
        self.assets_dir = self.project_root / "assets"
        self._check_assets_dir()

        # 2. 文件路径（需确保与实际文件名一致）
        self.xlsx_path = self.assets_dir / "BNA_subregions.xlsx"
        self.sc_matrix_path = self.assets_dir / "BNA_matrix_binary_246x246.csv"
        self.sc_4d_path = self.assets_dir / "BNA_SC_4D.nii.gz"
        self.fc_4d_path = self.assets_dir / "BNA_FC_4D.nii.gz"
        self.atlas_template_path = self.assets_dir / "BN_Atlas_246_2mm.nii.gz"  # 2mm模板

        # 3. 核心数据存储
        self.internal_name_to_id = {}  # 内部编号→ID
        self.id_to_internal_name = {}  # ID→内部编号
        self.anatomy_cn_mapping = {}  # 中文解剖描述→(中文, 内部编号)
        self.cn_anatomy_mapping = {}  # 中文解剖描述→内部编号
        self.sc_matrix = None  # 结构连接矩阵
        self.sc_4d_data = None  # 4D结构连接数据
        self.fc_4d_data = None  # 4D功能连接数据
        self.atlas_data = None  # 脑区模板数据
        self.excel_df = None  # Excel完整数据

        # 4. 加载流程（先Excel→解剖映射→其他数据）
        self.load_excel_data()
        self._init_anatomy_cn_mapping()
        self.load_region_mapping()
        self.load_sc_data()
        self.load_fc_data()
        self.load_atlas_template()
        self._verify_data()

    def _check_assets_dir(self):
        if not self.assets_dir.exists():
            raise FileNotFoundError(f"assets目录不存在！路径：{self.assets_dir}")

    def load_excel_data(self):
        if not self.xlsx_path.exists():
            raise FileNotFoundError(f"脑区映射表不存在：{self.xlsx_path}")
        self.excel_df = pd.read_excel(self.xlsx_path)
        required_cols = ["Left and Right Hemisphere", "Label ID.L", "Label ID.R",
                         "Anatomical and modified Cyto-architectonic descriptions"]
        for col in required_cols:
            if col not in self.excel_df.columns:
                raise KeyError(f"Excel缺少必要列！需包含：{required_cols}")
        self.excel_df["Anatomical and modified Cyto-architectonic descriptions"] = \
            self.excel_df["Anatomical and modified Cyto-architectonic descriptions"].fillna("未标注解剖描述")
        print("✅ Excel数据加载成功")

    def _init_anatomy_cn_mapping(self):
        """构建中文解剖描述映射（核心人机交互数据）"""
        for _, row in self.excel_df.iterrows():
            internal_name = row["Left and Right Hemisphere"].strip()
            anatomy_en = row["Anatomical and modified Cyto-architectonic descriptions"].strip()
            # 英文解剖描述→中文翻译（覆盖所有脑区类型）
            if "Superior Frontal Gyrus" in anatomy_en or "SFG" in internal_name:
                anatomy_cn = "上额回"
            elif "Middle Frontal Gyrus" in anatomy_en or "MFG" in internal_name:
                anatomy_cn = "中额回"
            elif "Inferior Frontal Gyrus" in anatomy_en or "IFG" in internal_name:
                anatomy_cn = "下额回"
            elif "Orbital Gyrus" in anatomy_en or "OrG" in internal_name:
                anatomy_cn = "眶额回"
            elif "Precentral Gyrus" in anatomy_en or "PrG" in internal_name:
                anatomy_cn = "中央前回"
            elif "Paracentral Lobule" in anatomy_en or "PCL" in internal_name:
                anatomy_cn = "中央旁小叶"
            elif "Superior Temporal Gyrus" in anatomy_en or "STG" in internal_name:
                anatomy_cn = "上颞回"
            elif "Middle Temporal Gyrus" in anatomy_en or "MTG" in internal_name:
                anatomy_cn = "中颞回"
            elif "Inferior Temporal Gyrus" in anatomy_en or "ITG" in internal_name:
                anatomy_cn = "下颞回"
            elif "Fusiform Gyrus" in anatomy_en or "FuG" in internal_name:
                anatomy_cn = "梭状回"
            elif "Parahippocampal Gyrus" in anatomy_en or "PhG" in internal_name:
                anatomy_cn = "海马旁回"
            elif "posterior Superior Temporal Sulcus" in anatomy_en or "pSTS" in internal_name:
                anatomy_cn = "上颞沟后部"
            elif "Superior Parietal Lobule" in anatomy_en or "SPL" in internal_name:
                anatomy_cn = "顶上小叶"
            elif "Inferior Parietal Lobule" in anatomy_en or "IPL" in internal_name:
                anatomy_cn = "顶下小叶"
            elif "Precuneus" in anatomy_en or "PCun" in internal_name:
                anatomy_cn = "楔前叶"
            elif "Postcentral Gyrus" in anatomy_en or "PoG" in internal_name:
                anatomy_cn = "中央后回"
            elif "Insular Gyrus" in anatomy_en or "INS" in internal_name:
                anatomy_cn = "岛叶回"
            elif "Cingulate Gyrus" in anatomy_en or "CG" in internal_name:
                anatomy_cn = "扣带回"
            elif "MedioVentral Occipital Cortex" in anatomy_en or "MVOcC" in internal_name:
                anatomy_cn = "腹内侧枕叶皮层"
            elif "lateral Occipital Cortex" in anatomy_en or "LOcC" in internal_name:
                anatomy_cn = "外侧枕叶皮层"
            elif "Amygdala" in anatomy_en or "Amyg" in internal_name:
                anatomy_cn = "杏仁核"
            elif "Hippocampus" in anatomy_en or "Hipp" in internal_name:
                anatomy_cn = "海马"
            elif "Basal Ganglia" in anatomy_en or "BG" in internal_name:
                anatomy_cn = "基底节"
            elif "Thalamus" in anatomy_en or "Tha" in internal_name:
                anatomy_cn = "丘脑"
            else:
                anatomy_cn = anatomy_en  # 未匹配时保留英文

            # 左/右脑映射构建
            left_full_cn = f"{anatomy_cn}（左脑）"
            self.anatomy_cn_mapping[left_full_cn] = (left_full_cn, internal_name)
            self.cn_anatomy_mapping[left_full_cn] = internal_name

            right_internal = f"{internal_name}_Right"
            right_full_cn = f"{anatomy_cn}（右脑）"
            self.anatomy_cn_mapping[right_full_cn] = (right_full_cn, right_internal)
            self.cn_anatomy_mapping[right_full_cn] = right_internal

        print(f"✅ 解剖描述映射完成（共{len(self.anatomy_cn_mapping)}个脑区）")

    def load_region_mapping(self):
        """内部编号-ID映射（后台查询用）"""
        for _, row in self.excel_df.iterrows():
            internal = row["Left and Right Hemisphere"].strip()
            left_id = row["Label ID.L"]
            right_id = row["Label ID.R"]

            if left_id not in self.id_to_internal_name:
                self.id_to_internal_name[left_id] = internal
            if internal not in self.internal_name_to_id:
                self.internal_name_to_id[internal] = left_id

            right_internal = f"{internal}_Right"
            if right_id not in self.id_to_internal_name:
                self.id_to_internal_name[right_id] = right_internal
            if right_internal not in self.internal_name_to_id:
                self.internal_name_to_id[right_internal] = right_id
        print(f"✅ 内部编号-ID映射完成（共{len(self.internal_name_to_id)}个）")

    def load_sc_data(self):
        for path in [self.sc_matrix_path, self.sc_4d_path]:
            if not path.exists():
                raise FileNotFoundError(f"结构连接文件不存在：{path}")
        self.sc_matrix = pd.read_csv(self.sc_matrix_path, header=None).values
        if self.sc_matrix.shape != (246, 246):
            raise ValueError(f"SC矩阵维度错误！需(246,246)，当前：{self.sc_matrix.shape}")
        sc_4d_img = nib.load(self.sc_4d_path)
        self.sc_4d_data = sc_4d_img.get_fdata()
        print("✅ 结构连接数据加载成功")

    def load_fc_data(self):
        if not self.fc_4d_path.exists():
            raise FileNotFoundError(f"功能连接文件不存在：{self.fc_4d_path}")
        fc_4d_img = nib.load(self.fc_4d_path)
        self.fc_4d_data = fc_4d_img.get_fdata()
        print("✅ 功能连接数据加载成功")

    def load_atlas_template(self):
        if not self.atlas_template_path.exists():
            raise FileNotFoundError(f"脑区模板不存在：{self.atlas_template_path}")
        atlas_img = nib.load(self.atlas_template_path)
        self.atlas_data = atlas_img.get_fdata()
        print("✅ 脑区模板加载成功")

    def _verify_data(self):
        """验证数据一致性"""
        if self.atlas_data.shape != self.fc_4d_data.shape[:3]:
            raise ValueError(f"维度不匹配！模板：{self.atlas_data.shape} | FC：{self.fc_4d_data.shape[:3]}")
        if self.atlas_data.shape != self.sc_4d_data.shape[:3]:
            raise ValueError(f"模板与SC数据维度不匹配！")
        if len(self.anatomy_cn_mapping) != 48:
            print(f"⚠️  解剖映射数量异常：当前{len(self.anatomy_cn_mapping)}个，预期48个")
        else:
            print("✅ 解剖映射数量正常")
        print("✅ 所有数据验证通过")

    def get_connections(self, input_anatomy_cn, threshold=0.1):
        """根据中文解剖描述查询连接结果"""
        # 1. 中文→内部编号→ID
        if input_anatomy_cn not in self.cn_anatomy_mapping:
            raise ValueError(f"未找到脑区：{input_anatomy_cn}")
        internal = self.cn_anatomy_mapping[input_anatomy_cn]
        if internal not in self.internal_name_to_id:
            raise ValueError(f"内部编号{internal}无对应ID")
        input_id = self.internal_name_to_id[internal]
        input_idx = input_id - 1

        # 2. 提取结构连接
        sc_ids = [i + 1 for i in range(246) if self.sc_matrix[input_idx, i] == 1 and i + 1 != input_id]
        sc_result = {}
        for rid in sc_ids:
            rid_internal = self.id_to_internal_name[rid]
            for cn, intl in self.cn_anatomy_mapping.items():
                if intl == rid_internal:
                    sc_result[cn] = 1
                    break

        # 3. 提取功能连接
        fc_result = {}
        for target_id in range(1, 247):
            if target_id == input_id:
                continue
            voxels = (self.atlas_data == target_id)
            if np.sum(voxels) == 0:
                continue
            mean_fc = np.mean(self.fc_4d_data[voxels, input_idx])
            if mean_fc >= threshold:
                target_internal = self.id_to_internal_name[target_id]
                for cn, intl in self.cn_anatomy_mapping.items():
                    if intl == target_internal:
                        fc_result[cn] = round(mean_fc, 3)
                        break

        return sc_result, fc_result


# --------------------------
# 2. PyQt5图形化界面类（核心优化：搜索框手动调高度+向上紧凑吸附）
# --------------------------
class BrainConnectionGUI(QMainWindow):
    def __init__(self, brain_tool):
        super().__init__()
        self.brain_tool = brain_tool
        self.current_suggestions = []  # 搜索结果列表数据
        self.init_ui()  # 初始化界面
        self.update_suggestions("")  # 初始加载完整脑区列表

    def init_ui(self):
        """初始化界面（按要求优化搜索框/结果框交互）"""
        # 窗口基础配置
        self.setWindowTitle("脑区结构/功能连接查询工具")
        self.setGeometry(100, 100, 400, 600)
        self.setMinimumSize(300, 300)

        # 中心部件与主布局（垂直布局，核心是让搜索区向上吸附）
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        main_layout = QVBoxLayout(central_widget)
        main_layout.setSpacing(15)
        main_layout.setContentsMargins(20, 20, 20, 20)

        # --------------------------
        # 1. 标题标签
        # --------------------------
        title_label = QLabel("脑区结构/功能连接查询")
        title_font = QFont()
        title_font.setPointSize(16)
        title_font.setBold(True)
        title_label.setFont(title_font)
        title_label.setAlignment(Qt.AlignCenter)
        main_layout.addWidget(title_label)

        # --------------------------
        # 2. 搜索核心区（用Widget包裹，实现向上紧凑吸附）
        # 关键：包裹后仅让该区域内部紧凑，不随窗口拉伸向下偏移
        # --------------------------
        search_core_widget = QWidget()  # 容器Widget，用于固定搜索区的向上对齐属性
        search_core_layout = QVBoxLayout(search_core_widget)  # 搜索区内部布局
        search_core_layout.setSpacing(8)  # 内部组件间距缩小，更紧凑
        search_core_layout.setContentsMargins(0, 0, 0, 0)  # 清除容器边距

        # 2.1 模糊查找标签
        input_label = QLabel("模糊查找")
        input_font = QFont()
        input_font.setPointSize(11)
        input_label.setFont(input_font)
        search_core_layout.addWidget(input_label)

        # 2.2 搜索输入框（光标常驻）
        self.search_edit = QLineEdit()
        search_font = QFont()
        search_font.setPointSize(12)
        self.search_edit.setFont(search_font)
        self.search_edit.setPlaceholderText("示例：额、上额、上额回、...")
        self.search_edit.textChanged.connect(self.update_suggestions)  # 输入实时更新结果
        self.search_edit.setFocus()  # 初始焦点在输入框
        search_core_layout.addWidget(self.search_edit)

        # 2.3 搜索结果标签
        result_box_label = QLabel("搜索结果")
        result_box_label.setFont(input_font)
        search_core_layout.addWidget(result_box_label)

        # 2.4 搜索结果窗格（核心优化：支持手动调高度，取消固定尺寸）
        self.suggest_list = QListWidget()
        suggest_font = QFont()
        suggest_font.setPointSize(11)
        self.suggest_list.setFont(suggest_font)
        # 关键设置：垂直方向可伸缩（支持鼠标拖动调整高度），启用滚动条
        self.suggest_list.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)
        self.suggest_list.setVerticalScrollBarPolicy(Qt.ScrollBarAsNeeded)  # 内容超出时显示滚动条
        self.suggest_list.itemDoubleClicked.connect(self.run_query)  # 双击触发查询
        search_core_layout.addWidget(self.suggest_list)

        # 将搜索核心区添加到主布局，设置对齐方式为“向上”
        main_layout.addWidget(search_core_widget, alignment=Qt.AlignTop)

        # --------------------------
        # 3. 控制区域（输出格式+查询按钮）
        # --------------------------
        control_layout = QHBoxLayout()

        # 输出格式选择（仅中文解剖描述）
        self.output_mode_combo = QComboBox()
        combo_font = QFont()
        combo_font.setPointSize(11)
        self.output_mode_combo.setFont(combo_font)
        self.output_mode_combo.addItems(["仅中文解剖描述"])
        self.output_mode_combo.setCurrentText("仅中文解剖描述")
        control_layout.addWidget(QLabel("输出格式："))
        control_layout.addWidget(self.output_mode_combo)
        control_layout.addStretch(1)

        # 查询按钮（始终启用，无选中项时提示）
        self.query_btn = QPushButton("查询")
        btn_font = QFont()
        btn_font.setPointSize(12)
        self.query_btn.setFont(btn_font)
        self.query_btn.clicked.connect(self.run_query)
        control_layout.addWidget(self.query_btn)

        main_layout.addLayout(control_layout)

        # --------------------------
        # 4. 分析结果展示框（占满剩余空间，随窗口拉伸）
        # --------------------------
        result_label = QLabel("查询结果：")
        result_label.setFont(input_font)
        main_layout.addWidget(result_label)

        self.result_text = QTextEdit()
        text_font = QFont()
        text_font.setPointSize(11)
        self.result_text.setFont(text_font)
        self.result_text.setReadOnly(True)
        self.result_text.setLineWrapMode(QTextEdit.WidgetWidth)
        # 关键：让结果框占满窗口剩余空间，实现“搜索区向上，结果区向下拉伸”的布局
        main_layout.addWidget(self.result_text, stretch=1)
        self.result_text.setText(f"双击/选中后回车/点击查询按钮，查询结构连接与静息态功能连接")

    def update_suggestions(self, input_text):
        """
        输入框文本变化时更新搜索结果（空输入显示所有脑区，常驻不隐藏）
        """
        input_text = input_text.lower().strip()
        self.current_suggestions = []

        # 空输入：显示所有中文解剖描述；有输入：模糊匹配
        if not input_text:
            self.current_suggestions = sorted(self.brain_tool.anatomy_cn_mapping.keys())  # 排序便于查找
        else:
            for anatomy_cn in self.brain_tool.anatomy_cn_mapping.keys():
                if input_text in anatomy_cn.lower():
                    self.current_suggestions.append(anatomy_cn)

        # 更新搜索结果框内容
        self.suggest_list.clear()
        for cn in self.current_suggestions:
            self.suggest_list.addItem(cn)

        # 始终选中第一项（便于快速查询）
        if len(self.current_suggestions) > 0:
            self.suggest_list.setCurrentRow(0)

        # 保持输入框焦点（即使更新结果后）
        self.search_edit.setFocus()

    def keyPressEvent(self, event: QKeyEvent):
        """
        键盘事件优化：
        1. 上下键控制搜索结果选中项，光标始终在输入框
        2. 回车键触发查询（无论焦点在哪）
        """
        # 处理上下键（控制搜索结果选中，不移动输入框光标）
        if len(self.current_suggestions) > 0:
            current_row = self.suggest_list.currentRow()
            total_rows = len(self.current_suggestions)

            # 下键：选中下一项（循环）
            if event.key() == Qt.Key_Down:
                new_row = (current_row + 1) % total_rows
                self.suggest_list.setCurrentRow(new_row)
                event.accept()  # 阻止事件传递（避免输入框光标下移）
                self.search_edit.setFocus()  # 强制回到输入框焦点
                return

            # 上键：选中上一项（循环）
            elif event.key() == Qt.Key_Up:
                new_row = (current_row - 1) % total_rows
                self.suggest_list.setCurrentRow(new_row)
                event.accept()
                self.search_edit.setFocus()
                return

        # 回车键触发查询（无论焦点在哪）
        if event.key() == Qt.Key_Return or event.key() == Qt.Key_Enter:
            self.run_query()
            event.accept()
            self.search_edit.setFocus()  # 查询后回到输入框
            return

        # 其他键（如字母、删除键）正常传递给输入框
        super().keyPressEvent(event)
        # 确保输入框始终有焦点（防止意外失去焦点）
        self.search_edit.setFocus()

    def run_query(self):
        """
        执行查询（优化点：查询前显示加载提示，仅用搜索结果选中项）
        """
        # 1. 获取选中的脑区（唯一查询来源）
        selected_item = self.suggest_list.currentItem()
        threshold = 0.1

        # 防护：无选中项时提示
        if not selected_item:
            self.result_text.setText("❌ 请先在搜索结果框中选择一个脑区！")
            self.search_edit.setFocus()
            return

        # 2. 获取选中项文本，显示加载提示
        input_anatomy_cn = selected_item.text()
        self.result_text.setText(f"🔍 正在查询 '{input_anatomy_cn}'...\n（请稍候）")
        # 强制刷新界面，确保加载提示显示（避免卡顿导致提示不出现）
        QApplication.processEvents()

        # 3. 调用工具查询连接
        try:
            sc_result, fc_result = self.brain_tool.get_connections(input_anatomy_cn, threshold)

            # 4. 格式化结果文本
            result = f"'{input_anatomy_cn}'的连接性：\n\n"

            # 结构连接结果
            result += f"1. 结构连接（共{len(sc_result)}个脑区）：\n"
            if sc_result:
                for i, (cn, _) in enumerate(sorted(sc_result.items()), 1):
                    result += f"   {i:2d}. {cn}\n"
            else:
                result += "   未找到相关结构连接脑区\n"

            # 功能连接结果
            result += f"\n2. 功能连接（强度≥{threshold}，共{len(fc_result)}个脑区）：\n"
            if fc_result:
                sorted_fc = sorted(fc_result.items(), key=lambda x: -x[1])
                for i, (cn, strength) in enumerate(sorted_fc, 1):
                    result += f"   {i:2d}. {cn} | 连接强度：{strength}\n"
            else:
                result += "   未找到符合阈值的功能连接脑区\n"

            # 显示最终结果
            self.result_text.setText(result)

        except ValueError as e:
            self.result_text.setText(f"❌ 查询失败：{str(e)}")
        except Exception as e:
            self.result_text.setText(f"❌ 程序异常：{str(e)}")

        # 查询后回到输入框焦点
        self.search_edit.setFocus()


# --------------------------
# 3. 程序入口
# --------------------------
if __name__ == "__main__":
    try:
        print("=== 开始加载脑区数据 ===")
        brain_tool = BrainConnectionTool()
        print("=== 所有数据加载完成，启动界面 ===")

        app = QApplication(sys.argv)
        gui = BrainConnectionGUI(brain_tool)
        gui.show()
        sys.exit(app.exec_())

    except Exception as e:
        print(f"❌ 程序启动失败：{e}")
        sys.exit(1)